from langchain.tools import BaseTool
from langchain.callbacks.manager import CallbackManagerForToolRun
from typing import Optional, Dict, Any, Type
from pydantic import BaseModel, Field
import json
import sys
from .parameter_collector import FloodForecastParameterCollector
from loguru import logger

# 配置loguru日志
logger.remove()  # 移除默认配置

# 控制台日志 - 显示INFO及以上级别
logger.add(
    sys.stdout,
    level="INFO",
    format="<green>{time:YYYY-MM-DD HH:mm:ss}</green> | <level>{level: <8}</level> | <cyan>{name}</cyan>:<cyan>{function}</cyan>:<cyan>{line}</cyan> - <level>{message}</level>",
    colorize=True
)

# 文件日志 - 详细日志，包含DEBUG级别
logger.add(
    "logs/flood_forecast_langchain_tool.log",
    level="DEBUG",
    format="{time:YYYY-MM-DD HH:mm:ss.SSS} | {level: <8} | {name}:{function}:{line} | {process.id} | {thread.id} - {message}",
    rotation="10 MB",
    retention="7 days",
    compression="zip",
    encoding="utf-8"
)

# 错误日志单独文件
logger.add(
    "logs/flood_forecast_langchain_tool_error.log",
    level="ERROR",
    format="{time:YYYY-MM-DD HH:mm:ss.SSS} | {level: <8} | {name}:{function}:{line} | {process.id} | {thread.id} - {message}\n{exception}",
    rotation="5 MB",
    retention="30 days",
    compression="zip",
    encoding="utf-8"
)

class FloodForecastParameterInput(BaseModel):
    """洪水预报参数收集工具的输入模式"""
    user_input: str = Field(description="用户关于洪水预报的输入或问题")
    session_id: Optional[str] = Field(default=None, description="会话ID，用于维持对话状态")
    action: str = Field(default="collect", description="操作类型：collect(收集参数), status(查看状态), reset(重置), final(获取最终JSON)")

class FloodForecastParameterTool(BaseTool):
    """洪水预报参数收集LangChain工具
    
    这个工具可以集成到LangChain的Agent中，用于收集洪水预报所需的参数
    """
    
    name: str = "flood_forecast_parameter_collector"
    description: str = """
    洪水预报参数收集工具。用于通过对话收集洪水预报所需的完整参数信息。
    
    功能：
    - 智能提取用户输入中的预报参数（水库名称、模型类型、时间等）
    - 针对缺失参数进行友好的反问
    - 维持对话状态，支持多轮交互
    - 生成最终的参数JSON供后续预测模型使用
    
    适用场景：
    - 用户提出洪水预报需求但信息不完整
    - 需要收集预报开始时间、预报时长、降雨预见期等关键参数
    - 为后续调用预测模型做准备
    
    输入格式：
    - user_input: 用户的输入文本
    - session_id: 可选的会话ID
    - action: 操作类型（collect/status/reset/final）
    """
    
    args_schema: Type[BaseModel] = FloodForecastParameterInput
    
    # 存储会话状态
    _sessions: Dict[str, FloodForecastParameterCollector] = {}
    
    def _run(
        self, 
        user_input: str, 
        session_id: Optional[str] = None,
        action: str = "collect",
        run_manager: Optional[CallbackManagerForToolRun] = None
    ) -> str:
        """执行工具"""
        try:
            # 生成会话ID
            if not session_id:
                import uuid
                session_id = str(uuid.uuid4())
            
            # 获取或创建收集器
            if session_id not in self._sessions:
                self._sessions[session_id] = FloodForecastParameterCollector()
            
            collector = self._sessions[session_id]
            
            # 根据操作类型执行不同逻辑
            if action == "collect":
                result = collector.process_user_input(user_input)
                
                # 构建返回信息
                response_parts = [
                    f"AI回复: {result['response']}",
                    f"\n会话ID: {session_id}",
                    f"状态: {result['status']}"
                ]
                
                if result['extracted_this_turn']:
                    response_parts.append(f"本轮提取的参数: {result['extracted_this_turn']}")
                
                collected_params = [k for k, v in result['parameters'].items() if v]
                if collected_params:
                    response_parts.append(f"已收集参数: {collected_params}")
                
                if result['missing_parameters']:
                    response_parts.append(f"缺失参数: {result['missing_parameters']}")
                
                if result['status'] == 'completed':
                    response_parts.append("\n✅ 参数收集完成！可以使用action='final'获取最终JSON")
                
                return "\n".join(response_parts)
                
            elif action == "status":
                status_info = {
                    "session_id": session_id,
                    "status": collector.conversation_state,
                    "parameters": {k: v["value"] for k, v in collector.required_parameters.items()},
                    "missing_parameters": collector._get_missing_parameters()
                }
                return f"会话状态: {json.dumps(status_info, ensure_ascii=False, indent=2)}"
                
            elif action == "reset":
                collector.reset()
                return f"会话 {session_id} 已重置"
                
            elif action == "final":
                if collector.conversation_state != "completed":
                    return "错误: 参数收集尚未完成，无法生成最终JSON"
                
                final_json = collector.get_final_parameters_json()
                return f"最终参数JSON:\n{final_json}"
                
            else:
                return f"错误: 不支持的操作类型 '{action}'"
                
        except Exception as e:
            return f"工具执行出错: {str(e)}"
    
    def get_session_info(self, session_id: str) -> Optional[Dict[str, Any]]:
        """获取会话信息"""
        if session_id not in self._sessions:
            return None
            
        collector = self._sessions[session_id]
        return {
            "session_id": session_id,
            "status": collector.conversation_state,
            "parameters": {k: v["value"] for k, v in collector.required_parameters.items()},
            "missing_parameters": collector._get_missing_parameters()
        }
    
    def cleanup_session(self, session_id: str) -> bool:
        """清理会话"""
        if session_id in self._sessions:
            del self._sessions[session_id]
            return True
        return False
    
    def list_active_sessions(self) -> list:
        """列出活跃会话"""
        return list(self._sessions.keys())

# 创建工具实例
flood_forecast_tool = FloodForecastParameterTool()

# LangChain Agent集成示例
def create_flood_forecast_agent():
    """创建包含洪水预报参数收集功能的Agent"""
    from langchain.agents import AgentExecutor, create_tool_calling_agent
    from langchain_core.prompts import ChatPromptTemplate
    from langchain_ollama import ChatOllama
    
    # 导入配置
    try:
        from config.config import OLLAMA_CONFIG
    except ImportError:
        # 如果无法导入配置，使用默认配置
        OLLAMA_CONFIG = {
            "base_url": "http://localhost:11434",
            "chat_model": "qwen:latest"
        }
    
    # 创建LLM
    llm = ChatOllama(
        base_url=OLLAMA_CONFIG["base_url"],
        model=OLLAMA_CONFIG["chat_model"],
        temperature=0.1
    )
    
    # 定义工具列表
    tools = [flood_forecast_tool]
    
    # 创建提示模板
    prompt = ChatPromptTemplate.from_messages([
        ("system", """
你是一个专业的洪水预报助手。你可以使用洪水预报参数收集工具来帮助用户收集预报所需的参数。

当用户提出洪水预报需求时：
1. 使用flood_forecast_parameter_collector工具收集参数
2. 如果参数不完整，工具会自动引导用户补充
3. 参数收集完成后，可以获取最终的JSON格式参数
4. 最终参数可以用于调用预测模型（通过MCP等方式）

请始终保持专业、友好的态度，帮助用户完成洪水预报参数的收集。
        """),
        ("human", "{input}"),
        ("placeholder", "{agent_scratchpad}")
    ])
    
    # 创建Agent
    agent = create_tool_calling_agent(llm, tools, prompt)
    
    # 创建Agent执行器
    agent_executor = AgentExecutor(
        agent=agent,
        tools=tools,
        verbose=True,
        handle_parsing_errors=True
    )
    
    return agent_executor

# 使用示例
if __name__ == "__main__":
    # 直接使用工具
    tool = FloodForecastParameterTool()
    
    # 测试参数收集
    result1 = tool._run(
        user_input="我要做一个大伙房水库的新安江预报，预测未来流量",
        session_id="test_session_1"
    )
    print("第一轮结果:")
    print(result1)
    print("\n" + "="*50 + "\n")
    
    # 继续收集参数
    result2 = tool._run(
        user_input="从明天上午8点开始，预报48小时",
        session_id="test_session_1"
    )
    print("第二轮结果:")
    print(result2)
    print("\n" + "="*50 + "\n")
    
    # 查看状态
    status = tool._run(
        user_input="",
        session_id="test_session_1",
        action="status"
    )
    print("状态查询:")
    print(status)